import os
import configparser
from dataclasses import dataclass
from project_path import get_project_path
from sqlalchemy import create_engine, text, inspect
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import sessionmaker, scoped_session
from contextlib import contextmanager
from urllib.parse import quote_plus


# 实例化配置文件
config = configparser.ConfigParser()


@dataclass
class DBConnect:
    def __init__(self):
        """获取数据库配置"""
        with open(os.path.join(get_project_path(), 'config.ini'), 'r', encoding='utf-8') as config_file:
            config.read_file(config_file)
        self.dialect = config['DATABASE']['dialect']
        self.host = config['DATABASE']['host']
        self.port = config['DATABASE']['port']
        self.user = config['DATABASE']['user']
        self.password = quote_plus(config['DATABASE']['password'])
        self.database = config['DATABASE']['database']

    def __post_init__(self):
        """创建数据库连接"""
        if "postgres" in self.dialect:
            self.driver = "postgresql+psycopg2"
        elif "mysql" in self.dialect:
            self.driver = "mysql+pymysql"
        self.engine = create_engine(
            "{}://{}:{}@{}:{}/{}".format(self.driver, self.user, self.password, self.host, self.port, self.database),
            pool_pre_ping=True,  # 检查连接是否可用
            pool_recycle=3600,  # 连接超时时间
            pool_size=5,  # 连接池大小
            max_overflow=10  # 连接池溢出时，最多允许创建的连接数
        )

    @contextmanager
    def get_session(self):
        """上下文管理数据库会话"""
        session = scoped_session(sessionmaker(bind=self.engine))
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()

    def connection(self):
        with self.get_session() as session:
            yield session

    def execute_sql(self, sql):
        """执行 SQL 查询"""
        try:
            with self.get_session() as session:
                sql_result = session.execute(text(sql))
            if sql_result.returns_rows:  # 若存在结果
                col_name = list(sql_result.keys())
                rows = sql_result.fetchall()
                # 处理列数据转置
                col_data = ([list(col) for col in zip(*rows)] if rows else [[] for _ in col_name])
                sql_result = dict(zip(col_name, col_data))
                return {'success': sql_result}
            else:
                return {'fail': "目前暂无相关数据"}
        except SQLAlchemyError as se:
            # 返回数据库错误信息
            return {'fail': str(se)}
        except Exception as e:
            # 返回其他类型的错误信息
            return {'fail': str(e)}

    def schema_info(self, insp, table_name):
        """数据库表结构"""
        # 获取表名及其注释
        schema_info = {'table_name': table_name}
        table_comment = insp.get_table_comment(table_name).get('text', '')
        if table_comment and table_comment != table_name:
            schema_info['table_comment'] = table_comment
        # 获取列信息
        columns = []
        for column in insp.get_columns(table_name):
            column_dict = {'name': column['name'], 'type': str(column['type'])}
            # 获取列注释
            if column['comment'] and column['comment'] != column['name']:
                column_dict['comment'] = column['comment']
            # 获取样例数据
            sql_result = self.execute_sql(f"SELECT {column['name']} FROM {table_name} LIMIT 5")
            if sql_result.get('success', ''):
                sample_data = list(sql_result.values())[0].get(column['name'], '')
                if sample_data != [None, None, None, None, None] and len(str(sample_data)) <= 200:  # 样例数据长度超过200，不予展示
                    column_dict['examples'] = [str(item) for item in sample_data]
            columns.append(column_dict)
        schema_info["columns"] = columns
        return schema_info

    def get_schema(self, table_name=None):
        """获取数据库表结构"""
        insp = inspect(self.engine)
        if table_name:
            return [self.schema_info(insp, table_name)]
        else:
            # for 循环遍历表名，添加 schema_info
            schema_info = []
            for table_name in insp.get_table_names():
                schema_info.append(self.schema_info(insp, table_name))
            return schema_info
